package com.gzist.springailab07.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.gzist.springailab07.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.*;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.swing.text.html.Option;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Deque;
import java.util.LinkedList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 处理工具调用的基础代理类，具体实现了 think 和 act 方法，可以用作创建实例的父类
 */
@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent {

    // 当前 SSE 输出器
    private transient SseEmitter currentEmitter;

    // 可用的工具
    private final ToolCallback[] availableTools;

    // 保存工具调用信息的响应结果（要调用那些工具）
    private ChatResponse toolCallChatResponse;

    // 工具调用管理者
    private final ToolCallingManager toolCallingManager;

    // 禁用 Spring AI 内置的工具调用机制，自己维护选项和消息上下文
    private final ChatOptions chatOptions;

    // ===== 补偿拦截器 start =====
    private final Deque<String> lastTools = new LinkedList<>();
    private int noToolCounter = 0;
    private static final int MAX_NO_TOOL = 3;   // 连续 3 次无工具 → 终止
    private static final int MAX_SAME_TOOL = 4;  // 同一工具 4 次 → 终止
    // ===== 补偿拦截器 end =====

    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        // 禁用 Spring AI 内置的工具调用机制，自己维护选项和消息上下文
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();
    }

    /**
     * 处理当前状态并决定下一步行动
     *
     * @return 是否需要执行行动
     */
    @Override
    public boolean think() {
        // 1、校验提示词，拼接用户提示词
        if (StrUtil.isNotBlank(getNextStepPrompt())) {
            getMessageList().add(new UserMessage(getNextStepPrompt()));
        }
        // 2、调用 AI 大模型，获取工具调用结果
        List<Message> messageList = getMessageList();
        Prompt prompt = new Prompt(messageList, this.chatOptions);
        try {
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();
            // 记录响应，用于等下 Act
            this.toolCallChatResponse = chatResponse;
            // 3、解析工具调用结果，获取要调用的工具
            // 助手消息
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            // 获取要调用的工具列表
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            // 输出提示信息
            String result = assistantMessage.getText();
            log.info("{}的思考：{}", getName(), result);
            log.info("{}选择了 {} 个工具来使用", getName(), toolCallList.size());
            // 将推理结果输出到 SSE
            if (StrUtil.isBlank(result)) {
                // 当决策工具使用的时候，推理内容为空
                result = "正在决策工具使用...";
            }
            safeSend(result);
            String toolCallInfo = toolCallList.stream()
                    .map(toolCall -> String.format("工具名称：%s，参数：%s", toolCall.name(), toolCall.arguments()))
                    .collect(Collectors.joining("\n"));
            log.info(toolCallInfo);

            // ===== 补偿拦截器 start =====
            updateCounters(toolCallList);
            if (shouldTerminateByCompensation(toolCallList)) {
                setState(AgentState.FINISHED);
                summary();
                return false;
            }
            // ===== 补偿拦截器 end =====

            // 如果不需要调用工具，返回 false
            if (toolCallList.isEmpty()) {
                // 只有不调用工具时，才需要手动记录助手消息
                getMessageList().add(assistantMessage);
                return false;
            } else {
                // 需要调用工具时，无需记录助手消息，因为调用工具时会自动记录
                return true;
            }
        } catch (Exception e) {
            log.error("{} 思考异常", getName(), e);
            safeSend("思考异常：" + e.getMessage());
            return false;
        }
    }

    /**
     * 执行工具调用并处理结果
     *
     * @return 执行结果
     */
    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()) {
            return "没有工具需要调用";
        }
        // 调用工具
        Prompt prompt = new Prompt(getMessageList(), this.chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        // 记录消息上下文，conversationHistory 已经包含了助手消息和工具调用返回的结果
        setMessageList(toolExecutionResult.conversationHistory());
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        // 判断是否调用了终止工具
        boolean terminateToolCalled = toolResponseMessage.getResponses().stream()
                .anyMatch(response -> response.name().equals("doTerminate"));

        if (terminateToolCalled) {
            // 任务结束，更改状态
            setState(AgentState.FINISHED);
            summary();
            return "任务结束！";
        }
        List<AssistantMessage.ToolCall> toolCalls = toolCallChatResponse
                .getResult().getOutput().getToolCalls();
        List<ToolResponseMessage.ToolResponse> responses = toolResponseMessage.getResponses();

        // 推送工具调用结果
        String message = buildToolMessage(toolCalls, responses);
        safeSend(message);
        return "工具调用成功！";
    }

    /**
     * 任务最终回复推送
     */
    private void summary() {
        String result = "";
        List<Message> history = new ArrayList<>(getMessageList());
        history.add(new SystemMessage(getFinalSystemPrompt()));
        history.add(new UserMessage(getFinalUserPrompt()));
        ChatResponse summaryResp = getChatClient()
                .prompt(new Prompt(history, chatOptions))
                .system(getSystemPrompt())
                .call()
                .chatResponse();
        if (summaryResp != null) {
            result = summaryResp.getResult().getOutput().getText();
        }
        safeSend(result);
        safeSend("<END>");
    }

    /**
     * 构建工具调用结果消息体
     *
     * @param calls     工具调用列表
     * @param responses 工具调用结果列表
     * @return 工具调用结果消息
     */
    private static String buildToolMessage(List<AssistantMessage.ToolCall> calls,
                                           List<ToolResponseMessage.ToolResponse> responses) {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < calls.size(); i++) {
            AssistantMessage.ToolCall call = calls.get(i);
            ToolResponseMessage.ToolResponse resp = responses.get(i);
            sb.append(String.format(
                    "🔧 工具调用\n名称：%s\n参数：%s\n结果：%s\n",
                    call.name(),
                    call.arguments(),
                    resp.responseData()
            ));
        }

        return sb.toString();
    }

    /**
     * 安全推送内容到sse服务端
     *
     * @param text 要发送的文本
     */
    private void safeSend(String text) {
        if (currentEmitter == null) return;
        try {
            currentEmitter.send(text);
        } catch (IOException e) {
            log.warn("SSE 推送失败", e);
        }
    }

    private void updateCounters(List<AssistantMessage.ToolCall> toolCalls) {
        if (toolCalls.isEmpty()) {
            noToolCounter++;
            lastTools.offerLast("NO_TOOL");
        } else {
            noToolCounter = 0;
            toolCalls.forEach(c -> lastTools.offerLast(c.name()));
        }
        while (lastTools.size() > Math.max(MAX_NO_TOOL, MAX_SAME_TOOL)) {
            lastTools.pollFirst();
        }
    }

    private boolean shouldTerminateByCompensation(List<AssistantMessage.ToolCall> toolCalls) {
        // 连续无工具
        if (noToolCounter >= MAX_NO_TOOL) return true;
        // 连续同一工具
        if (lastTools.size() < MAX_SAME_TOOL) return false;
        String last = lastTools.peekLast();
        return lastTools.stream()
                .skip(lastTools.size() - MAX_SAME_TOOL)
                .allMatch(last::equals);
    }
}